import json
import os
import sys
import time
import threading
import queue
import keyboard
import pyaudio
import edge_tts
import asyncio
import pygame
from vosk import Model, KaldiRecognizer
import requests
import tempfile
import unicodedata

# 从环境变量中获取API密钥
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_KEY", "sk-bab1e4f901344202b2f514efe1200128")  # 使用默认值作为备选

# 初始化pygame音频模块
pygame.mixer.init()

class DeepseekAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.api_url = "https://api.deepseek.com/v1/chat/completions"
        self.conversation_history = []
    
    def clean_markdown(self, text):
        """清理文本中的markdown语法和特殊符号"""
        import re
        
        # 移除代码块
        text = re.sub(r'```[\s\S]*?```', '', text)
        
        # 移除内联代码
        text = re.sub(r'`([^`]*)`', r'\1', text)
        
        # 移除标题符号
        text = re.sub(r'^#+\s+', '', text, flags=re.MULTILINE)
        
        # 移除粗体和斜体
        text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)
        text = re.sub(r'\*(.*?)\*', r'\1', text)
        text = re.sub(r'__(.*?)__', r'\1', text)
        text = re.sub(r'_(.*?)_', r'\1', text)
        
        # 移除删除线
        text = re.sub(r'~~(.*?)~~', r'\1', text)
        
        # 移除链接
        text = re.sub(r'\[(.*?)\]\(.*?\)', r'\1', text)
        
        # 移除图片
        text = re.sub(r'!\[(.*?)\]\(.*?\)', r'\1', text)
        
        # 移除引用
        text = re.sub(r'^>\s+', '', text, flags=re.MULTILINE)
        
        # 移除HTML标签
        text = re.sub(r'<[^>]*>', '', text)
        
        # 移除任务列表
        text = re.sub(r'^- \[[xX ]\]\s+', '- ', text, flags=re.MULTILINE)
        
        # 移除水平线
        text = re.sub(r'^-{3,}$|^_{3,}$|^\*{3,}$', '', text, flags=re.MULTILINE)
        
        # 移除表格结构
        text = re.sub(r'\|.*\|', '', text)
        text = re.sub(r'^\s*[-:]+[-|\s:]*$', '', text, flags=re.MULTILINE)
        
        # 移除表情符号和特殊字符
        # 1. 移除 :smile: 格式的表情符号
        text = re.sub(r':[a-zA-Z_]+:', '', text)
        
        # 2. 使用unicodedata来过滤表情符号
        clean_text = ""
        for char in text:
            # 检查每个字符的类别，排除"So"(符号)类别的字符
            if 'So' not in unicodedata.category(char):
                clean_text += char
        text = clean_text
        
        # 处理markdown列表
        text = re.sub(r'^\s*[*+-]\s+', '- ', text, flags=re.MULTILINE)  # 无序列表
        text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)  # 有序列表
        
        # 处理换行
        text = re.sub(r'\n{3,}', '\n\n', text)  # 最多保留两个连续空行
        
        return text.strip()
    
    def ask(self, question):
        self.conversation_history.append({"role": "user", "content": question})
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        data = {
            "model": "deepseek-chat",
            "messages": self.conversation_history,
            "stream": True
        }
        
        response = requests.post(
            self.api_url,
            headers=headers,
            json=data,
            stream=True
        )
        
        answer = ""
        current_line = ""
        for line in response.iter_lines():
            if line:
                line_text = line.decode('utf-8')
                if line_text.startswith("data: "):
                    json_str = line_text[6:]
                    if json_str == "[DONE]":
                        break
                    try:
                        chunk_data = json.loads(json_str)
                        if "choices" in chunk_data and chunk_data["choices"]:
                            content = chunk_data["choices"][0].get("delta", {}).get("content", "")
                            if content:
                                # 添加到总回复
                                answer += content
                                
                                # 累积当前行的内容
                                current_line += content
                                
                                # 如果有换行符，清理当前行并输出
                                if '\n' in content:
                                    parts = current_line.split('\n')
                                    for i, part in enumerate(parts):
                                        if i < len(parts) - 1:  # 对除了最后一部分外的所有部分
                                            cleaned_part = self.clean_markdown(part)
                                            print(cleaned_part, end="\n", flush=True)
                                    
                                    # 保留最后一部分作为新的当前行
                                    current_line = parts[-1]
                                else:
                                    # 直接输出内容，无需处理
                                    print(content, end="", flush=True)
                    except json.JSONDecodeError:
                        continue
        
        # 处理最后可能剩余的内容
        if current_line:
            cleaned_line = self.clean_markdown(current_line)
            print(cleaned_line, flush=True)
        else:
            print()  # 换行
            
        # 清理整个回复并保存到对话历史
        cleaned_answer = self.clean_markdown(answer)
        self.conversation_history.append({"role": "assistant", "content": cleaned_answer})
        return cleaned_answer

class VoiceRecognizer:
    def __init__(self, model_path):
        self.model = Model(model_path)
        self.recognizer = KaldiRecognizer(self.model, 16000)
        self.audio = pyaudio.PyAudio()
        self.recording = False
        self.audio_queue = queue.Queue()
        self.result_text = ""
        self.audio_thread = None
        
    def _audio_callback(self):
        while self.recording:
            try:
                data = self.stream.read(4000, exception_on_overflow=False)
                if len(data) > 0 and self.recording:
                    if self.recognizer.AcceptWaveform(data):
                        # 处理中间结果，如果需要的话
                        pass
                    time.sleep(0.01)
            except Exception as e:
                print(f"录音错误: {e}")
                break
        
    def start_recording(self):
        if self.recording:  # 如果已经在录音，避免重复启动
            return
            
        self.recording = True
        self.stream = self.audio.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=16000,
            input=True,
            frames_per_buffer=4000
        )
        self.recognizer.Reset()
        print("🎤 正在录音...", end="", flush=True)
        
        # 启动录音线程
        self.audio_thread = threading.Thread(target=self._audio_callback)
        self.audio_thread.daemon = True
        self.audio_thread.start()
        
    def stop_recording(self):
        if not self.recording:
            return ""
            
        self.recording = False
        
        # 等待录音线程结束
        if self.audio_thread and self.audio_thread.is_alive():
            self.audio_thread.join(timeout=1.0)
            
        self.stream.stop_stream()
        self.stream.close()
        print(" 录音结束，正在识别...")
        
        # 获取最终识别结果
        result = json.loads(self.recognizer.FinalResult())
        self.result_text = result.get("text", "")
        print(f"🔊 识别结果: {self.result_text}")
        return self.result_text

class TextToSpeech:
    def __init__(self):
        self.voices = {
            "1": "zh-CN-XiaoxiaoNeural",
            "2": "zh-CN-YunxiNeural",
            "3": "zh-CN-YunjianNeural",
            "4": "zh-CN-XiaoyiNeural",
            "5": "zh-CN-YunyangNeural",
            "6": "zh-CN-liaoning-XiaobeiNeural",
            "7": "zh-TW-HsiaoChenNeural",
            "8": "zh-TW-YunJheNeural"
        }
        self.current_voice = self.voices["1"]
        self.speech_enabled = True
        self.is_playing = False
    
    def set_voice(self, voice_num):
        if voice_num in self.voices:
            self.current_voice = self.voices[voice_num]
            return True
        return False
    
    def toggle_speech(self):
        self.speech_enabled = not self.speech_enabled
        return self.speech_enabled
    
    def stop_speech(self):
        """停止当前正在播放的语音"""
        if self.is_playing and pygame.mixer.music.get_busy():
            pygame.mixer.music.stop()
            self.is_playing = False
            return True
        return False
    
    def list_voices(self):
        voice_list = ""
        for key, value in self.voices.items():
            voice_list += f"{key}: {value}\n"
        return voice_list
    
    async def speak(self, text):
        if not self.speech_enabled or not text:
            return
        
        self.is_playing = True
        communicate = edge_tts.Communicate(text, self.current_voice)
        
        # 创建临时文件
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
        temp_file.close()
        
        try:
            await communicate.save(temp_file.name)
            
            # 使用pygame播放音频
            pygame.mixer.music.load(temp_file.name)
            pygame.mixer.music.play()
            
            # 等待播放完成
            while pygame.mixer.music.get_busy() and self.is_playing:
                pygame.time.Clock().tick(10)
                
        finally:
            # 播放结束后重置状态
            self.is_playing = False
            # 删除临时文件
            try:
                os.unlink(temp_file.name)
            except:
                pass

class VoiceAssistant:
    def __init__(self, api_key, model_path):
        self.deepseek = DeepseekAPI(api_key)
        self.recognizer = VoiceRecognizer(model_path)
        self.tts = TextToSpeech()
        self.running = True
        self.current_input = ""  # 存储当前输入行的内容
        self.recording_active = False  # 记录录音状态
        
    def welcome(self):
        welcome_text = "您好！语音助手已就绪，随时为您效劳。"
        print(f"🤖 {welcome_text}")
        asyncio.run(self.tts.speak(welcome_text))
        
    def show_help(self):
        help_text = """
指令列表:
/help   - 显示帮助信息
/exit   - 退出程序
/voice  - 切换语音播放开关
/voice_list - 显示可用语音列表
/voice_set <编号> - 设置语音(1-8)
/stop   - 停止语音播放
空格键  - 在输入框为空时，按住空格键开始录音，松开识别
ESC键   - 按下ESC键可随时停止语音播放
        """
        print(help_text)
        return help_text
    
    def _keyboard_hook(self, event):
        # 使用keyboard模块处理所有键盘事件
        try:
            # 按ESC键停止语音播放
            if event.name == 'esc' and event.event_type == keyboard.KEY_DOWN and self.tts.is_playing:
                self.tts.stop_speech()
                print("\n🔇 已停止语音播放")
                return True
                
            # 按下空格键且未处于录音状态，开始录音
            # 但仅当当前输入为空时才触发录音功能
            if (event.name == 'space' and 
                event.event_type == keyboard.KEY_DOWN and 
                not self.recording_active and 
                not self.recognizer.recording and 
                not self.current_input.strip()):  # 当前输入为空时才触发录音
                
                self.recording_active = True
                self.recognizer.start_recording()
                return False  # 阻止空格键事件传递
                
            # 释放空格键且处于录音状态，停止录音
            elif event.name == 'space' and event.event_type == keyboard.KEY_UP and self.recording_active:
                self.recording_active = False
                user_input = self.recognizer.stop_recording()
                if user_input:
                    # 保存当前输入
                    current_text = self.current_input
                    # 处理识别结果
                    self.process_input(user_input)
                    # 恢复输入提示和文本
                    print("\n请输入问题或指令(输入框为空时按住空格键可语音输入): ", end="", flush=True)
                    print(current_text, end="", flush=True)
                    self.current_input = current_text
                return False  # 阻止空格键事件传递
                
            # 跟踪当前输入的文本
            if event.event_type == keyboard.KEY_DOWN and not self.recording_active:
                if event.name == 'space':
                    self.current_input += ' '
                elif event.name == 'backspace' and self.current_input:
                    self.current_input = self.current_input[:-1]
                elif event.name == 'enter':
                    # 处理回车键，暂不记录，由input函数处理
                    pass
                elif len(event.name) == 1:  # 普通字符
                    self.current_input += event.name
                    
        except Exception as e:
            print(f"\n键盘处理错误: {e}")
            
        return True  # 允许其他按键正常传递
        
    def process_input(self, user_input):
        if not user_input:
            return
            
        if user_input.lower() in ["/exit", "退出"]:
            self.exit()
            return
            
        if user_input.lower() in ["/help", "help", "帮助"]:
            self.show_help()
            return
            
        if user_input.lower() == "/voice":
            enabled = self.tts.toggle_speech()
            status = "开启" if enabled else "关闭"
            print(f"🔊 语音播放已{status}")
            return
            
        if user_input.lower() == "/voice_list":
            voice_list = self.tts.list_voices()
            print(f"可用语音列表:\n{voice_list}")
            return
            
        if user_input.lower().startswith("/voice_set "):
            try:
                voice_num = user_input.split()[1]
                if self.tts.set_voice(voice_num):
                    print(f"语音已设置为: {self.tts.current_voice}")
                else:
                    print("无效的语音编号，请使用1-8")
            except:
                print("设置语音失败，格式: /voice_set <编号>")
            return
            
        if user_input.lower() in ["/stop", "停止"]:
            if self.tts.stop_speech():
                print("🔇 已停止语音播放")
            else:
                print("当前没有语音在播放")
            return
            
        # 发送到Deepseek处理
        print(f"\n🧠 思考中...\n")
        response = self.deepseek.ask(user_input)
        
        # 语音输出 - response已经在ask方法中被清理过
        asyncio.run(self.tts.speak(response))
    
    def start_keyboard_listener(self):
        # 使用keyboard模块注册钩子
        keyboard.hook(self._keyboard_hook)
        
    def run(self):
        self.welcome()
        self.show_help()
        self.start_keyboard_listener()
        
        try:
            while self.running:
                self.current_input = ""  # 重置当前输入
                user_input = input("\n请输入问题或指令(输入框为空时按住空格键可语音输入): ").strip()
                if user_input:
                    self.process_input(user_input)
        except KeyboardInterrupt:
            pass
        finally:
            self.exit()
    
    def exit(self):
        self.running = False
        # 取消所有键盘钩子
        keyboard.unhook_all()
        print("再见，感谢使用!")
        sys.exit(0)

if __name__ == "__main__":
    MODEL_PATH = "vosk-model-cn-0.22"
    
    # 检查模型路径是否存在
    if not os.path.exists(MODEL_PATH):
        print(f"错误: 找不到语音识别模型，请确认'{MODEL_PATH}'目录存在")
        sys.exit(1)
    
    assistant = VoiceAssistant(DEEPSEEK_API_KEY, MODEL_PATH)
    assistant.run() 